import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import json
import sys
import os
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def main(imputed_csv_path, target_samples):
    try:
        print(f"开始噪声扩充，数据文件: {imputed_csv_path}")
        print(f"目标样本数: {target_samples}")

        # 读取插补后的数据
        df = pd.read_csv(imputed_csv_path)
        original_samples = len(df)

        print(f"原始插补后数据形状: {df.shape}")
        print(f"数据列: {list(df.columns)}")

        # 完全按照你的Jupyter代码方法 - 分成10个split
        print("将数据分成10个split...")
        splits = np.array_split(df, 10)

        print(f"每个split的大小:")
        for i, split in enumerate(splits):
            print(f"  Split {i+1}: {len(split)} 行")

        # 存储扩充后的数据
        augmented_data = []
        current_samples = 0
        splits_used = 0
        noise_applications = 0

        # 对前7个split添加噪声进行扩充
        print("\n开始对前7个split进行噪声扩充...")

        split_index = 0
        while current_samples < target_samples and split_index < 7:
            print(f"\n处理 Split {split_index + 1}...")

            split_data = splits[split_index].copy()

            # 添加高斯噪声
            noise_level = 0.01  # 噪声强度
            noisy_split = add_gaussian_noise(split_data, noise_level)

            augmented_data.append(noisy_split)
            current_samples += len(noisy_split)
            noise_applications += 1

            print(f"  添加噪声后样本数: {len(noisy_split)}")
            print(f"  累计样本数: {current_samples}")

            # 如果当前split的数据加完后还没达到目标，继续使用同一个split
            if current_samples < target_samples:
                # 继续使用同一个split，添加不同强度的噪声
                continue
            else:
                split_index += 1

        splits_used = split_index + 1

        print(f"\n总共使用了 {splits_used} 个split，应用噪声 {noise_applications} 次")

        # 合并所有扩充的数据
        if augmented_data:
            final_df = pd.concat(augmented_data, ignore_index=True)
            # 确保正好是目标数量
            if len(final_df) > target_samples:
                final_df = final_df.head(target_samples)
        else:
            final_df = df.head(target_samples)

        final_samples = len(final_df)
        augmentation_ratio = final_samples / original_samples if original_samples > 0 else 0

        print(f"\n扩充完成!")
        print(f"原始样本数: {original_samples}")
        print(f"最终样本数: {final_samples}")
        print(f"扩充倍数: {augmentation_ratio:.2f}x")

        # 保存扩充后的数据
        os.makedirs('temp/output', exist_ok=True)
        output_path = 'temp/output/augmented_data.csv'
        final_df.to_csv(output_path, index=False)

        print(f"扩充后数据已保存到: {output_path}")

        # 创建扩充效果可视化
        visualization_paths = create_augmentation_visualization(df, final_df)

        # 计算数据质量指标
        quality_analysis = analyze_augmentation_quality(df, final_df)

        # 构建结果
        result = {
            'augmented_data_path': output_path,
            'original_samples': int(original_samples),
            'final_samples': int(final_samples),
            'target_samples': int(target_samples),
            'augmentation_ratio': float(augmentation_ratio),
            'splits_used': int(splits_used),
            'noise_applications': int(noise_applications),
            'noise_level_used': 0.01,
            'columns': list(final_df.columns),
            'data_preview': final_df.head(5).to_dict('records'),
            'visualization_paths': visualization_paths,
            'quality_analysis': quality_analysis,
            'status': 'success'
        }

        print(json.dumps(result))

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'噪声扩充失败: {str(e)}'
        }
        print(json.dumps(error_result))

def add_gaussian_noise(data, noise_level):
    """为数据添加高斯噪声"""
    noisy_data = data.copy()

    for col in data.select_dtypes(include=[np.number]).columns:
        # 计算该列的标准差来确定噪声强度
        col_std = data[col].std()
        if not np.isnan(col_std) and col_std > 0:
            # 生成高斯噪声
            noise = np.random.normal(0, col_std * noise_level, len(data))
            noisy_data[col] = data[col] + noise

            # 对于特定字段，确保值在合理范围内
            if col == 'qyb':  # 气油比不应该为负数
                noisy_data[col] = np.maximum(noisy_data[col], 0)

    return noisy_data

def create_augmentation_visualization(original_df, augmented_df):
    """创建扩充效果可视化"""
    visualization_paths = {}

    try:
        # 1. 样本数量对比
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
        categories = ['原始数据', '扩充后数据']
        counts = [len(original_df), len(augmented_df)]

        bars = ax.bar(categories, counts, color=['skyblue', 'lightcoral'])
        ax.set_ylabel('样本数量')
        ax.set_title('数据扩充前后样本数量对比')

        # 在柱子上显示数值
        for bar, count in zip(bars, counts):
            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(counts)*0.01,
                    f'{count:,}', ha='center', va='bottom')

        plt.tight_layout()
        sample_count_path = 'temp/output/sample_count_comparison.png'
        plt.savefig(sample_count_path, dpi=300, bbox_inches='tight')
        plt.close()
        visualization_paths['sample_count'] = sample_count_path

        # 2. 数据分布对比（选择几个关键字段）
        key_fields = ['qyb'] + [col for col in original_df.columns if col != 'qyb'][:5]

        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.flatten()

        for i, field in enumerate(key_fields):
            if i >= 6:
                break

            ax = axes[i]

            # 原始数据分布
            ax.hist(original_df[field].dropna(), bins=30, alpha=0.5,
                    label='原始数据', color='blue', density=True)

            # 扩充后数据分布
            ax.hist(augmented_df[field].dropna(), bins=30, alpha=0.5,
                    label='扩充后数据', color='red', density=True)

            ax.set_title(f'{field} 分布对比')
            ax.legend()
            ax.grid(True, alpha=0.3)

        # 隐藏多余的子图
        for i in range(len(key_fields), 6):
            axes[i].set_visible(False)

        plt.suptitle('扩充前后数据分布对比', fontsize=16, fontweight='bold')
        plt.tight_layout()

        distribution_path = 'temp/output/distribution_comparison.png'
        plt.savefig(distribution_path, dpi=300, bbox_inches='tight')
        plt.close()
        visualization_paths['distribution'] = distribution_path

    except Exception as e:
        print(f"创建可视化失败: {e}")

    return visualization_paths

def analyze_augmentation_quality(original_df, augmented_df):
    """分析扩充质量"""
    quality_analysis = {}

    try:
        # 统计信息对比
        quality_analysis['statistical_comparison'] = {}

        for col in original_df.select_dtypes(include=[np.number]).columns:
            original_stats = original_df[col].describe()
            augmented_stats = augmented_df[col].describe()

            quality_analysis['statistical_comparison'][col] = {
                'mean_original': float(original_stats['mean']),
                'mean_augmented': float(augmented_stats['mean']),
                'std_original': float(original_stats['std']),
                'std_augmented': float(augmented_stats['std']),
                'mean_change_percent': float((augmented_stats['mean'] - original_stats['mean']) / original_stats['mean'] * 100) if original_stats['mean'] != 0 else 0,
                'std_change_percent': float((augmented_stats['std'] - original_stats['std']) / original_stats['std'] * 100) if original_stats['std'] != 0 else 0
            }

        # 整体质量评估
        quality_analysis['overall_quality'] = {
            'data_consistency': 'good',  # 可以基于统计变化程度来评估
            'noise_level': 'appropriate',
            'sample_increase': len(augmented_df) / len(original_df)
        }

    except Exception as e:
        print(f"质量分析失败: {e}")
        quality_analysis = {'error': str(e)}

    return quality_analysis

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(json.dumps({
            'status': 'error',
            'message': '参数错误，需要提供CSV文件路径'
        }))
    else:
        target = int(sys.argv[2]) if len(sys.argv) > 2 else 40000
        main(sys.argv[1], target)
